{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8b1b12fd",
   "metadata": {},
   "source": [
    "# Benchmarking FastKafka app"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e79e80d6",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "To benchmark a `FastKafka` project, you will need the following:\n",
    "\n",
    "1. A library built with `FastKafka`.\n",
    "2. A running `Kafka` instance to benchmark the FastKafka application against."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36a63ab0",
   "metadata": {},
   "source": [
    "### Creating FastKafka Code\n",
    "\n",
    "Let's create a `FastKafka`-based application and write it to the `application.py` file based on the [tutorial](/docs#tutorial).\n",
    "\n",
    "```python\n",
    "# content of the \"application.py\" file\n",
    "\n",
    "from contextlib import asynccontextmanager\n",
    "\n",
    "from sklearn.datasets import load_iris\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "\n",
    "from fastkafka import FastKafka\n",
    "\n",
    "ml_models = {}\n",
    "\n",
    "\n",
    "@asynccontextmanager\n",
    "async def lifespan(app: FastKafka):\n",
    "    # Load the ML model\n",
    "    X, y = load_iris(return_X_y=True)\n",
    "    ml_models[\"iris_predictor\"] = LogisticRegression(random_state=0, max_iter=500).fit(\n",
    "        X, y\n",
    "    )\n",
    "    yield\n",
    "    # Clean up the ML models and release the resources\n",
    "    ml_models.clear()\n",
    "\n",
    "\n",
    "from pydantic import BaseModel, NonNegativeFloat, Field\n",
    "\n",
    "class IrisInputData(BaseModel):\n",
    "    sepal_length: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Sepal length in cm\"\n",
    "    )\n",
    "    sepal_width: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Sepal width in cm\"\n",
    "    )\n",
    "    petal_length: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Petal length in cm\"\n",
    "    )\n",
    "    petal_width: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Petal width in cm\"\n",
    "    )\n",
    "\n",
    "\n",
    "class IrisPrediction(BaseModel):\n",
    "    species: str = Field(..., example=\"setosa\", description=\"Predicted species\")\n",
    "\n",
    "from fastkafka import FastKafka\n",
    "\n",
    "kafka_brokers = {\n",
    "    \"localhost\": {\n",
    "        \"url\": \"localhost\",\n",
    "        \"description\": \"local development kafka broker\",\n",
    "        \"port\": 9092,\n",
    "    },\n",
    "    \"production\": {\n",
    "        \"url\": \"kafka.airt.ai\",\n",
    "        \"description\": \"production kafka broker\",\n",
    "        \"port\": 9092,\n",
    "        \"protocol\": \"kafka-secure\",\n",
    "        \"security\": {\"type\": \"plain\"},\n",
    "    },\n",
    "}\n",
    "\n",
    "kafka_app = FastKafka(\n",
    "    title=\"Iris predictions\",\n",
    "    kafka_brokers=kafka_brokers,\n",
    "    lifespan=lifespan,\n",
    ")\n",
    "\n",
    "@kafka_app.consumes(topic=\"input_data\", auto_offset_reset=\"latest\")\n",
    "async def on_input_data(msg: IrisInputData):\n",
    "    species_class = ml_models[\"iris_predictor\"].predict(\n",
    "        [[msg.sepal_length, msg.sepal_width, msg.petal_length, msg.petal_width]]\n",
    "    )[0]\n",
    "\n",
    "    await to_predictions(species_class)\n",
    "\n",
    "\n",
    "@kafka_app.produces(topic=\"predictions\")\n",
    "async def to_predictions(species_class: int) -> IrisPrediction:\n",
    "    iris_species = [\"setosa\", \"versicolor\", \"virginica\"]\n",
    "\n",
    "    prediction = IrisPrediction(species=iris_species[species_class])\n",
    "    return prediction\n",
    "\n",
    "```\n",
    "\n",
    "\n",
    "`FastKafka` has a decorator for benchmarking which is appropriately called as `benchmark`.\n",
    "Let's edit our `application.py` file and add the `benchmark` decorator to the consumes method.\n",
    "\n",
    "\n",
    "```python\n",
    "# content of the \"application.py\" file with benchmark\n",
    "\n",
    "from contextlib import asynccontextmanager\n",
    "\n",
    "from sklearn.datasets import load_iris\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "\n",
    "from fastkafka import FastKafka\n",
    "\n",
    "ml_models = {}\n",
    "\n",
    "\n",
    "@asynccontextmanager\n",
    "async def lifespan(app: FastKafka):\n",
    "    # Load the ML model\n",
    "    X, y = load_iris(return_X_y=True)\n",
    "    ml_models[\"iris_predictor\"] = LogisticRegression(random_state=0, max_iter=500).fit(\n",
    "        X, y\n",
    "    )\n",
    "    yield\n",
    "    # Clean up the ML models and release the resources\n",
    "    ml_models.clear()\n",
    "\n",
    "\n",
    "from pydantic import BaseModel, NonNegativeFloat, Field\n",
    "\n",
    "class IrisInputData(BaseModel):\n",
    "    sepal_length: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Sepal length in cm\"\n",
    "    )\n",
    "    sepal_width: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Sepal width in cm\"\n",
    "    )\n",
    "    petal_length: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Petal length in cm\"\n",
    "    )\n",
    "    petal_width: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Petal width in cm\"\n",
    "    )\n",
    "\n",
    "\n",
    "class IrisPrediction(BaseModel):\n",
    "    species: str = Field(..., example=\"setosa\", description=\"Predicted species\")\n",
    "\n",
    "from fastkafka import FastKafka\n",
    "\n",
    "kafka_brokers = {\n",
    "    \"localhost\": {\n",
    "        \"url\": \"localhost\",\n",
    "        \"description\": \"local development kafka broker\",\n",
    "        \"port\": 9092,\n",
    "    },\n",
    "    \"production\": {\n",
    "        \"url\": \"kafka.airt.ai\",\n",
    "        \"description\": \"production kafka broker\",\n",
    "        \"port\": 9092,\n",
    "        \"protocol\": \"kafka-secure\",\n",
    "        \"security\": {\"type\": \"plain\"},\n",
    "    },\n",
    "}\n",
    "\n",
    "kafka_app = FastKafka(\n",
    "    title=\"Iris predictions\",\n",
    "    kafka_brokers=kafka_brokers,\n",
    "    lifespan=lifespan,\n",
    ")\n",
    "\n",
    "@kafka_app.consumes(topic=\"input_data\", auto_offset_reset=\"latest\")\n",
    "@kafka_app.benchmark(interval=1, sliding_window_size=5)\n",
    "async def on_input_data(msg: IrisInputData):\n",
    "    species_class = ml_models[\"iris_predictor\"].predict(\n",
    "        [[msg.sepal_length, msg.sepal_width, msg.petal_length, msg.petal_width]]\n",
    "    )[0]\n",
    "\n",
    "    await to_predictions(species_class)\n",
    "\n",
    "\n",
    "@kafka_app.produces(topic=\"predictions\")\n",
    "async def to_predictions(species_class: int) -> IrisPrediction:\n",
    "    iris_species = [\"setosa\", \"versicolor\", \"virginica\"]\n",
    "\n",
    "    prediction = IrisPrediction(species=iris_species[species_class])\n",
    "    return prediction\n",
    "\n",
    "```\n",
    "\n",
    "\n",
    "Here we are conducting a benchmark of a function that consumes data from the `input_data` topic with an interval of 1 second and a sliding window size of 5. \n",
    "\n",
    "This `benchmark` method uses the `interval` parameter to calculate the results over a specific time period, and the `sliding_window_size` parameter to determine the maximum number of results to use in calculating the average throughput and standard deviation. \n",
    "\n",
    "This benchmark is important to ensure that the function is performing optimally and to identify any areas for improvement."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69cd27f1",
   "metadata": {},
   "source": [
    "### Starting Kafka\n",
    "\n",
    "If you already have a `Kafka` running somewhere, then you can skip this step. \n",
    "\n",
    "Please keep in mind that your benchmarking results may be affected by bottlenecks such as network, CPU cores in the Kafka machine, or even the Kafka configuration itself."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b147ee29",
   "metadata": {},
   "source": [
    "#### Installing Java and Kafka\n",
    "We need a working `Kafka`instance to benchmark our `FastKafka` app, and to run `Kafka` we need `Java`. Thankfully, `FastKafka` comes with a CLI to install both `Java` and `Kafka` on our machine.\n",
    "\n",
    "So, let's install `Java` and `Kafka` by executing the following command.\n",
    "\n",
    "```cmd\n",
    "fastkafka testing install_deps\n",
    "```\n",
    "\n",
    "The above command will extract `Kafka` scripts at the location \"$HOME/.local/kafka_2.13-3.3.2\" on your machine. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a460d8d",
   "metadata": {},
   "source": [
    "#### Creating configuration for Zookeeper and Kafka\n",
    "Now we need to start `Zookeeper` and `Kafka` separately, and to start them we need `zookeeper.properties` and `kafka.properties` files.\n",
    "\n",
    "Let's create a folder inside the folder where `Kafka` scripts were extracted and change directory into it.\n",
    "\n",
    "```cmd\n",
    "mkdir $HOME/.local/kafka_2.13-3.3.2/data_dir && cd $HOME/.local/kafka_2.13-3.3.2/data_dir\n",
    "```\n",
    "\n",
    "Let's create a file called `zookeeper.properties` and write the following content to the file:\n",
    "\n",
    "```txt\n",
    "dataDir=$HOME/.local/kafka_2.13-3.3.2/data_dir/zookeeper\n",
    "clientPort=2181\n",
    "maxClientCnxns=0\n",
    "```\n",
    "\n",
    "Similarly, let's create a file called `kafka.properties` and write the following content to the file:\n",
    "\n",
    "```txt\n",
    "broker.id=0\n",
    "listeners=PLAINTEXT://:9092\n",
    "\n",
    "num.network.threads=3\n",
    "num.io.threads=8\n",
    "socket.send.buffer.bytes=102400\n",
    "socket.receive.buffer.bytes=102400\n",
    "socket.request.max.bytes=104857600\n",
    "\n",
    "num.partitions=1\n",
    "num.recovery.threads.per.data.dir=1\n",
    "offsets.topic.replication.factor=1\n",
    "transaction.state.log.replication.factor=1\n",
    "transaction.state.log.min.isr=1\n",
    "\n",
    "log.dirs=$HOME/.local/kafka_2.13-3.3.2/data_dir/kafka_logs\n",
    "log.flush.interval.messages=10000\n",
    "log.flush.interval.ms=1000\n",
    "log.retention.hours=168\n",
    "log.retention.bytes=1073741824\n",
    "log.segment.bytes=1073741824\n",
    "log.retention.check.interval.ms=300000\n",
    "\n",
    "zookeeper.connect=localhost:2181\n",
    "zookeeper.connection.timeout.ms=18000\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07392f57",
   "metadata": {},
   "source": [
    "#### Starting Zookeeper and Kafka\n",
    "\n",
    "We need two different terminals to run `Zookeeper` in one and `Kafka` in another.\n",
    "Let's open a new terminal and run the following commands to start `Zookeeper`:\n",
    "\n",
    "```cmd\n",
    "export PATH=$PATH:$HOME/.jdk/jdk-11.0.18+10/bin\n",
    "cd $HOME/.local/kafka_2.13-3.3.2/bin\n",
    "./zookeeper-server-start.sh ../data_dir/zookeeper.properties\n",
    "```\n",
    "\n",
    "\n",
    "Once `Zookeeper` is up and running, open a new terminal and execute the follwing commands to start `Kafka`:\n",
    "```cmd\n",
    "export PATH=$PATH:$HOME/.jdk/jdk-11.0.18+10/bin\n",
    "cd $HOME/.local/kafka_2.13-3.3.2/bin\n",
    "./kafka-server-start.sh ../data_dir/kafka.properties\n",
    "```\n",
    "\n",
    "Now we have both `Zookeeper` and `Kafka` up and running."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac98a6ac",
   "metadata": {},
   "source": [
    "#### Creating topics in Kafka\n",
    "\n",
    "In a new terminal, please execute the following command to create necessary topics in `Kafka`:\n",
    "\n",
    "```cmd\n",
    "export PATH=$PATH:$HOME/.jdk/jdk-11.0.18+10/bin\n",
    "cd $HOME/.local/kafka_2.13-3.3.2/bin\n",
    "./kafka-topics.sh --create --topic input_data --partitions 6 --bootstrap-server localhost:9092\n",
    "./kafka-topics.sh --create --topic predictions --partitions 6 --bootstrap-server localhost:9092\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e947d67",
   "metadata": {},
   "source": [
    "\n",
    "#### Populating topics with dummy data\n",
    "\n",
    "To benchmark our `FastKafka` app, we need some data in `Kafka` topics. \n",
    "\n",
    "In the same terminal, let's create some dummy data:\n",
    "\n",
    "\n",
    "```cmd\n",
    "yes '{\"sepal_length\": 0.7739560486, \"sepal_width\": 0.8636615789, \"petal_length\": 0.6122663046, \"petal_width\": 0.1338914722}' | head -n 1000000 > /tmp/test_data\n",
    "```\n",
    "\n",
    "\n",
    "This command will create a file called `test_data` in the `tmp` folder with one million rows of text. This will act as dummy data to populate the `input_data` topic.\n",
    "\n",
    "Let's populate the created topic `input_data` with the dummy data which we created above:\n",
    "\n",
    "```cmd\n",
    "./kafka-console-producer.sh --bootstrap-server localhost:9092 --topic input_data < /tmp/test_data\n",
    "```\n",
    "\n",
    "Now our topic `input_data` has one million records/messages in it. If you want more messages in topic, you can simply execute the above command again and again."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fac21518",
   "metadata": {},
   "source": [
    "### Benchmarking FastKafka\n",
    "\n",
    "\n",
    "Once `Zookeeper` and `Kafka` are ready, benchmarking `FastKafka` app is as simple as running the `fastkafka run` command:\n",
    "\n",
    "```cmd\n",
    "fastkafka run --num-workers 1 --kafka-broker localhost application:kafka_app\n",
    "```\n",
    "\n",
    "This command will start the `FastKafka` app and begin consuming messages from `Kafka`, which we spun up earlier. \n",
    "Additionally, the same command will output all of the benchmark throughputs based on the `interval` and `sliding_window_size` values.\n",
    "\n",
    "The output for the `fastkafka run` command is:\n",
    "\n",
    "```txt\n",
    "[385814]: 23-04-07 10:49:18.380 [INFO] application: Current group id is ZDGTBVWVBBDMZCW\n",
    "[385814]: 23-04-07 10:49:18.382 [INFO] fastkafka._application.app: set_kafka_broker() : Setting bootstrap_servers value to 'localhost:9092'\n",
    "[385814]: 23-04-07 10:49:18.382 [INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
    "[385814]: 23-04-07 10:49:18.387 [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
    "[385814]: 23-04-07 10:49:18.387 [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'group_id': 'ZDGTBVWVBBDMZCW', 'auto_offset_reset': 'earliest', 'bootstrap_servers': 'localh\n",
    "ost:9092', 'max_poll_records': 100}\n",
    "[385814]: 23-04-07 10:49:18.390 [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
    "[385814]: 23-04-07 10:49:18.390 [INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'input_data'})\n",
    "[385814]: 23-04-07 10:49:18.390 [INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'input_data'}\n",
    "[385814]: 23-04-07 10:49:18.390 [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
    "[385814]: 23-04-07 10:49:18.395 [INFO] aiokafka.consumer.group_coordinator: Discovered coordinator 0 for group ZDGTBVWVBBDMZCW\n",
    "[385814]: 23-04-07 10:49:18.395 [INFO] aiokafka.consumer.group_coordinator: Revoking previously assigned partitions set() for group ZDGTBVWVBBDMZCW\n",
    "[385814]: 23-04-07 10:49:18.395 [INFO] aiokafka.consumer.group_coordinator: (Re-)joining group ZDGTBVWVBBDMZCW\n",
    "[385814]: 23-04-07 10:49:21.396 [INFO] aiokafka.consumer.group_coordinator: Joined group 'ZDGTBVWVBBDMZCW' (generation 1) with member_id aiokafka-0.8.0-b1f06560-6983-4d5e-a9af-8084e0e652cc\n",
    "[385814]: 23-04-07 10:49:21.396 [INFO] aiokafka.consumer.group_coordinator: Elected group leader -- performing partition assignments using roundrobin\n",
    "[385814]: 23-04-07 10:49:21.397 [INFO] aiokafka.consumer.group_coordinator: Successfully synced group ZDGTBVWVBBDMZCW with generation 1\n",
    "[385814]: 23-04-07 10:49:21.397 [INFO] aiokafka.consumer.group_coordinator: Setting newly assigned partitions {TopicPartition(topic='input_data', partition=0), TopicPartition(topic='input_data', partition=1), TopicPartition(topic='input_data', partition\n",
    "=2), TopicPartition(topic='input_data', partition=3)} for group ZDGTBVWVBBDMZCW\n",
    "[385814]: 23-04-07 10:49:22.409 [INFO] fastkafka.benchmark: Throughput = 93,598, Avg throughput = 93,598 - For application.on_input_data(interval=1,sliding_window_size=5)\n",
    "[385814]: 23-04-07 10:49:23.409 [INFO] fastkafka.benchmark: Throughput = 91,847, Avg throughput = 92,723 - For application.on_input_data(interval=1,sliding_window_size=5)\n",
    "[385814]: 23-04-07 10:49:24.409 [INFO] fastkafka.benchmark: Throughput = 92,948, Avg throughput = 92,798 - For application.on_input_data(interval=1,sliding_window_size=5)\n",
    "[385814]: 23-04-07 10:49:25.409 [INFO] fastkafka.benchmark: Throughput = 93,227, Avg throughput = 92,905 - For application.on_input_data(interval=1,sliding_window_size=5)\n",
    "[385814]: 23-04-07 10:49:26.409 [INFO] fastkafka.benchmark: Throughput = 93,553, Avg throughput = 93,035 - For application.on_input_data(interval=1,sliding_window_size=5)\n",
    "[385814]: 23-04-07 10:49:27.409 [INFO] fastkafka.benchmark: Throughput = 92,699, Avg throughput = 92,855 - For application.on_input_data(interval=1,sliding_window_size=5)\n",
    "[385814]: 23-04-07 10:49:28.409 [INFO] fastkafka.benchmark: Throughput = 92,716, Avg throughput = 93,029 - For application.on_input_data(interval=1,sliding_window_size=5)\n",
    "[385814]: 23-04-07 10:49:29.409 [INFO] fastkafka.benchmark: Throughput = 92,897, Avg throughput = 93,019 - For application.on_input_data(interval=1,sliding_window_size=5)\n",
    "[385814]: 23-04-07 10:49:30.409 [INFO] fastkafka.benchmark: Throughput = 92,854, Avg throughput = 92,944 - For application.on_input_data(interval=1,sliding_window_size=5)\n",
    "[385814]: 23-04-07 10:49:31.410 [INFO] fastkafka.benchmark: Throughput = 92,672, Avg throughput = 92,768 - For application.on_input_data(interval=1,sliding_window_size=5)\n",
    "```\n",
    "\n",
    "Based on the output, when using 1 worker, our `FastKafka` app achieved a `throughput` of 93k messages per second and an `average throughput` of 93k messages per second."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
